#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
PDF工具 - 一个跨平台的PDF处理CLI工具
支持PDF合并、拆分、压缩和信息查看功能
"""

import argparse
import sys
import os
from typing import List
from pypdf import PdfWriter, PdfReader


def merge_pdfs(input_files: List[str], output_file: str):
    """
    合并多个PDF文件为一个PDF文件
    
    Args:
        input_files: 输入的PDF文件列表
        output_file: 输出的PDF文件路径
    """
    # 创建PDF写入器
    writer = PdfWriter()
    
    # 遍历所有输入文件
    for input_file in input_files:
        if not os.path.exists(input_file):
            print(f"错误: 文件 '{input_file}' 不存在")
            return False
        
        try:
            # 读取PDF文件
            reader = PdfReader(input_file)
            
            # 将每一页添加到写入器中
            for page in reader.pages:
                writer.add_page(page)
                
            print(f"已添加文件: {input_file}")
        except Exception as e:
            print(f"错误: 读取文件 '{input_file}' 失败: {e}")
            return False
    
    # 写入合并后的PDF文件
    try:
        with open(output_file, "wb") as output_pdf:
            writer.write(output_pdf)
        print(f"成功合并PDF文件到: {output_file}")
        return True
    except Exception as e:
        print(f"错误: 写入文件 '{output_file}' 失败: {e}")
        return False


def split_pdf(input_file: str, page_range: str, output_dir: str):
    """
    根据页码范围拆分PDF文件
    
    Args:
        input_file: 输入的PDF文件路径
        page_range: 页码范围，如"1-3,5,-1"
        output_dir: 输出目录
    """
    # 检查输入文件是否存在
    if not os.path.exists(input_file):
        print(f"错误: 文件 '{input_file}' 不存在")
        return False
    
    # 创建输出目录
    if not os.path.exists(output_dir):
        try:
            os.makedirs(output_dir)
        except Exception as e:
            print(f"错误: 创建目录 '{output_dir}' 失败: {e}")
            return False
    
    try:
        # 读取PDF文件
        reader = PdfReader(input_file)
        total_pages = len(reader.pages)
        print(f"PDF文件总页数: {total_pages}")
        
        # 解析页码范围
        page_numbers = parse_page_range(page_range, total_pages)
        if not page_numbers:
            print("错误: 页码范围解析失败")
            return False
        
        print(f"需要拆分的页面: {page_numbers}")
        
        # 为每个页面创建单独的PDF文件
        for i, page_num in enumerate(page_numbers):
            # 创建新的PDF写入器
            writer = PdfWriter()
            
            # 添加指定页面
            page_index = page_num - 1  # 转换为0基索引
            if 0 <= page_index < total_pages:
                writer.add_page(reader.pages[page_index])
                
                # 生成输出文件名
                output_file = os.path.join(output_dir, f"page_{page_num}.pdf")
                
                # 写入PDF文件
                with open(output_file, "wb") as output_pdf:
                    writer.write(output_pdf)
                
                print(f"已创建页面 {page_num} 的PDF文件: {output_file}")
            else:
                print(f"警告: 页面 {page_num} 超出范围 (1-{total_pages})")
        
        print(f"成功拆分PDF文件到目录: {output_dir}")
        return True
    except Exception as e:
        print(f"错误: 拆分PDF文件失败: {e}")
        return False


def parse_page_range(range_str: str, total_pages: int) -> List[int]:
    """
    解析页码范围字符串
    
    Args:
        range_str: 页码范围字符串，如"1-3,5,-1"
        total_pages: 总页数
        
    Returns:
        页码列表
    """
    if not range_str:
        return []
    
    pages = set()
    ranges = range_str.split(',')
    
    for r in ranges:
        r = r.strip()
        if not r:
            continue
            
        if '-' in r:
            # 处理范围
            parts = r.split('-')
            if len(parts) == 2:
                start_str, end_str = parts
                
                # 处理起始页
                if start_str:
                    start = int(start_str)
                else:
                    start = 1  # 如果没有起始页，默认为1
                
                # 处理结束页
                if end_str:
                    if end_str == '1':  # 特殊处理-1
                        end = total_pages
                    else:
                        end = int(end_str)
                else:
                    end = total_pages  # 如果没有结束页，默认为最后一页
                
                # 添加范围内的所有页面
                for i in range(start, min(end + 1, total_pages + 1)):
                    pages.add(i)
        else:
            # 处理单个页面
            try:
                page_num = int(r)
                if page_num == -1:
                    pages.add(total_pages)
                elif 1 <= page_num <= total_pages:
                    pages.add(page_num)
            except ValueError:
                print(f"警告: 无法解析页码 '{r}'")
    
    return sorted(list(pages))


# 尝试导入PyMuPDF库（用于有损压缩）
try:
    import fitz  # PyMuPDF
    HAS_PYMUPDF = True
except ImportError:
    HAS_PYMUPDF = False
    print("警告: 未安装PyMuPDF库，有损压缩功能将不可用")


def compress_pdf(input_file: str, output_file: str, level: str, dpi: int = 150, quality: int = 70):
    """
    压缩PDF文件
    
    Args:
        input_file: 输入的PDF文件路径
        output_file: 输出的PDF文件路径
        level: 压缩等级 (lossless|medium|high)
        dpi: 图像DPI (默认150)
        quality: 图像质量 (默认70)
    """
    # 检查输入文件是否存在
    if not os.path.exists(input_file):
        print(f"错误: 文件 '{input_file}' 不存在")
        return False
    
    try:
        if level == 'lossless':
            # 无损压缩：使用pypdf压缩内容流、移除元数据/嵌入文件
            return _compress_lossless(input_file, output_file)
        elif level in ['medium', 'high']:
            # 有损压缩：如果安装了PyMuPDF，优先对图片重采样与有损重压
            if HAS_PYMUPDF:
                return _compress_lossy(input_file, output_file, dpi, quality)
            else:
                print(f"警告: 未安装PyMuPDF库，回退到无损压缩")
                return _compress_lossless(input_file, output_file)
        else:
            print(f"错误: 不支持的压缩等级 '{level}'")
            return False
    except Exception as e:
        print(f"错误: 压缩PDF文件失败: {e}")
        return False


def _compress_lossless(input_file: str, output_file: str) -> bool:
    """
    无损压缩PDF文件
    
    Args:
        input_file: 输入的PDF文件路径
        output_file: 输出的PDF文件路径
        
    Returns:
        压缩是否成功
    """
    try:
        # 读取PDF文件
        reader = PdfReader(input_file)
        
        # 创建PDF写入器
        writer = PdfWriter()
        
        # 复制所有页面
        for page in reader.pages:
            writer.add_page(page)
        
        # 压缩内容流
        writer.compress_identical_objects()
        
        # 移除元数据
        writer._info = None
        
        # 写入压缩后的PDF文件
        with open(output_file, "wb") as output_pdf:
            writer.write(output_pdf)
        
        # 显示文件大小对比
        input_size = os.path.getsize(input_file)
        output_size = os.path.getsize(output_file)
        compression_ratio = (1 - output_size / input_size) * 100
        
        print(f"成功无损压缩PDF文件: {input_file} -> {output_file}")
        print(f"文件大小: {input_size} bytes -> {output_size} bytes (压缩率: {compression_ratio:.2f}%)")
        
        return True
    except Exception as e:
        print(f"错误: 无损压缩PDF文件失败: {e}")
        return False


def _compress_lossy(input_file: str, output_file: str, dpi: int, quality: int) -> bool:
    """
    有损压缩PDF文件（使用PyMuPDF）
    
    Args:
        input_file: 输入的PDF文件路径
        output_file: 输出的PDF文件路径
        dpi: 图像DPI
        quality: 图像质量
        
    Returns:
        压缩是否成功
    """
    try:
        # 打开PDF文件
        doc = fitz.open(input_file)
        
        # 检查PDF是否包含图像
        has_images = False
        total_images = 0
        
        # 遍历所有页面检查是否有图像
        for page_num in range(doc.page_count):
            page = doc[page_num]
            image_list = page.get_images()
            if image_list:
                has_images = True
                total_images += len(image_list)
        
        # 如果PDF不包含图像，则回退到无损压缩
        if not has_images:
            print("提示: 此PDF文件不包含图像，使用无损压缩更合适")
            doc.close()
            return _compress_lossless(input_file, output_file)
        
        print(f"检测到PDF包含{total_images}张图像，使用有损压缩")
        
        # 创建新的PDF文档
        new_doc = fitz.open()
        
        # 遍历每一页
        for page_num in range(doc.page_count):
            page = doc[page_num]
            
            # 获取页面的像素图
            # 限制DPI以减少文件大小
            actual_dpi = min(dpi, 150)  # 限制最大DPI为150
            pix = page.get_pixmap(dpi=actual_dpi)
            
            # 创建新的PDF页面
            new_page = new_doc.new_page(width=page.rect.width, height=page.rect.height)
            
            # 将像素图转换为JPEG格式并插入到新页面
            # 限制质量参数以确保压缩效果
            actual_quality = max(10, min(quality, 95))  # 限制质量在10-95之间
            img_data = pix.pil_tobytes(format="JPEG", quality=actual_quality)
            new_page.insert_image(new_page.rect, stream=img_data)
            
            # 清理资源
            pix = None
        
        # 保存压缩后的PDF文件，使用压缩选项
        new_doc.save(output_file, garbage=4, deflate=True, clean=True)
        new_doc.close()
        doc.close()
        
        # 显示文件大小对比
        input_size = os.path.getsize(input_file)
        output_size = os.path.getsize(output_file)
        compression_ratio = (1 - output_size / input_size) * 100
        
        print(f"成功有损压缩PDF文件: {input_file} -> {output_file}")
        print(f"文件大小: {input_size} bytes -> {output_size} bytes (压缩率: {compression_ratio:.2f}%)")
        print(f"压缩参数: DPI={actual_dpi}, 质量={actual_quality}")
        
        return True
    except Exception as e:
        print(f"错误: 有损压缩PDF文件失败: {e}")
        return False


def show_pdf_info(input_file: str):
    """
    显示PDF文件信息
    
    Args:
        input_file: 输入的PDF文件路径
    """
    # 检查输入文件是否存在
    if not os.path.exists(input_file):
        print(f"错误: 文件 '{input_file}' 不存在")
        return False
    
    try:
        # 读取PDF文件
        reader = PdfReader(input_file)
        
        # 获取基本信息
        total_pages = len(reader.pages)
        print(f"文件名: {os.path.basename(input_file)}")
        print(f"总页数: {total_pages}")
        
        # 获取文档信息
        if reader.metadata:
            metadata = reader.metadata
            print("文档信息:")
            if metadata.title:
                print(f"  标题: {metadata.title}")
            if metadata.author:
                print(f"  作者: {metadata.author}")
            if metadata.creator:
                print(f"  创建者: {metadata.creator}")
            if metadata.producer:
                print(f"  生产者: {metadata.producer}")
            if metadata.subject:
                print(f"  主题: {metadata.subject}")
            if metadata.creation_date:
                print(f"  创建日期: {metadata.creation_date}")
            if metadata.modification_date:
                print(f"  修改日期: {metadata.modification_date}")
        else:
            print("文档信息: 无")
        
        # 获取第一页的信息作为示例
        if total_pages > 0:
            first_page = reader.pages[0]
            print("第一页信息:")
            print(f"  媒体框: {first_page.mediabox}")
            print(f"  裁剪框: {first_page.cropbox}")
            print(f"  旋转角度: {first_page.get('/Rotate', 0)}")
        
        return True
    except Exception as e:
        print(f"错误: 读取PDF文件信息失败: {e}")
        return False


def main():
    """主函数，处理命令行参数并执行相应操作"""
    parser = argparse.ArgumentParser(
        description="PDF工具 - 支持PDF合并、拆分、压缩和信息查看",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  pdf_tool.py merge -o out.pdf in1.pdf in2.pdf ...
  pdf_tool.py split in.pdf -r "1-3,5,-1" -o out_dir
  pdf_tool.py compress in.pdf -o out.pdf -l lossless|medium|high [--dpi 150] [--quality 70]
  pdf_tool.py info in.pdf
        """
    )
    
    # 创建子命令解析器
    subparsers = parser.add_subparsers(dest='command', help='可用的子命令')
    
    # merge子命令
    merge_parser = subparsers.add_parser('merge', help='合并多个PDF文件')
    merge_parser.add_argument('input_files', nargs='+', help='要合并的PDF文件')
    merge_parser.add_argument('-o', '--output', required=True, help='输出PDF文件路径')
    
    # split子命令
    split_parser = subparsers.add_parser('split', help='拆分PDF文件')
    split_parser.add_argument('input_file', help='要拆分的PDF文件')
    split_parser.add_argument('-r', '--range', required=True, help='页码范围，如"1-3,5,-1"')
    split_parser.add_argument('-o', '--output', required=True, help='输出目录')
    
    # compress子命令
    compress_parser = subparsers.add_parser('compress', help='压缩PDF文件')
    compress_parser.add_argument('input_file', help='要压缩的PDF文件')
    compress_parser.add_argument('-o', '--output', required=True, help='输出PDF文件路径')
    compress_parser.add_argument('-l', '--level', choices=['lossless', 'medium', 'high'], 
                                required=True, help='压缩等级')
    compress_parser.add_argument('--dpi', type=int, default=150, help='图像DPI (默认150)')
    compress_parser.add_argument('--quality', type=int, default=70, help='图像质量 (默认70)')
    
    # info子命令
    info_parser = subparsers.add_parser('info', help='查看PDF文件信息')
    info_parser.add_argument('input_file', help='要查看信息的PDF文件')
    
    # 解析命令行参数
    args = parser.parse_args()
    
    # 根据子命令执行相应操作
    if args.command == 'merge':
        merge_pdfs(args.input_files, args.output)
    elif args.command == 'split':
        split_pdf(args.input_file, args.range, args.output)
    elif args.command == 'compress':
        compress_pdf(args.input_file, args.output, args.level, args.dpi, args.quality)
    elif args.command == 'info':
        show_pdf_info(args.input_file)
    else:
        parser.print_help()


if __name__ == '__main__':
    main()